#ifndef SPLINECURVEFIT_HPP_
#define SPLINECURVEFIT_HPP_

#include "BSpline.hpp"
#include <algorithm>
#include <cmath>
#include <stdexcept>
#include <vector>

class SplineCurveFit {
    public:
        typedef std::vector<double> vector_type;
        typedef BSpline spline_type;
        enum knots_type{
            cumulative = 1,
            equaldistance =2
        };
        enum curve_type{
            regular = 1,
            ball = 2
        };
    private:
        int dimension=2;
        knots_type kn_type=cumulative;
        curve_type cv_type=regular;
        double theta=0;
        vector_type xData=vector_type();
        vector_type yData=vector_type();
        vector_type zData=vector_type();
        vector_type knots=vector_type();
        spline_type xCurve=spline_type();
        spline_type yCurve=spline_type();
        spline_type zCurve=spline_type();
        spline_type xBall=spline_type(3);
        spline_type yBall=spline_type(3);
    public:
        SplineCurveFit(int dim=2): dimension(dim){}
        //
        SplineCurveFit(vector_type x, vector_type y, int dim=2, knots_type type=cumulative):
            dimension(dim), xData(x), yData(y), kn_type(type) {}
        //
        SplineCurveFit(vector_type x, vector_type y, vector_type z, int dim=3, knots_type type=cumulative):
            dimension(dim), xData(x), yData(y), zData(z), kn_type(type) {}
        //
        ~SplineCurveFit(){}
        //
        void generate_knots(){
            checkData();
            int n = xData.size();
            knots.clear();
            knots.resize(n);
            knots[0] = 0;

            if (dimension==2) {
                if (kn_type == cumulative) {
                    for (int i=1; i<n; i++) {
                        knots[i] = sqrt(pow(xData[i]-xData[i-1],2)+pow(yData[i]-yData[i-1],2))+knots[i-1];
                    }
                    for (int i=1; i<n; i++) {
                        knots[i] = knots[i]/knots[n-1];
                    }
                } else {
                    for (int i=1 ; i<n; i++) {
                        knots[i] = 1.0*i/(n-1);
                    }
                }
            } else if (dimension==3) {
                if (kn_type == cumulative) {
                    for (int i=1; i<n; i++) {
                        knots[i] = sqrt(pow(xData[i]-xData[i-1],2)+pow(yData[i]-yData[i-1],2)+pow(zData[i]-zData[i-1],2))+knots[i-1];
                    }
                    for (int i=1; i<n; i++) {
                        knots[i] = knots[i]/knots[n-1];
                    }
                } else {
                    for (int i=1 ; i<n; i++) {
                        knots[i] = 1.0*i/n-1;
                    }
                }
            } else {
                throw std::invalid_argument("invalid dimension");
            }
        }
        //
        void setDimension(int dim){
            dimension = dim;
        }
        //
        void setData(vector_type x, vector_type y, vector_type z=vector_type()){
            xData = x;
            yData = y;
            zData = z;
        }
        //
        void setKnots(vector_type x){
            knots = x;
        }
        //
        void setKnotsType(knots_type type) {
            kn_type = type;
        }
        //
        void checkData(){
            if (dimension==2) {
                if (xData.size()!=yData.size()) {
                    throw std::invalid_argument("invalid data");
                }
            }else if (dimension==3) {
                if (xData.size()!=yData.size()||xData.size()!=zData.size()) {
                    throw std::invalid_argument("invalid data");
                }
            }else {
                throw std::invalid_argument("invalid dimension");
            }
        }
        //
        void setupCurve(){
            generate_knots();
            if (dimension==2) {
                xCurve.setData(knots, xData);
                xCurve.setBoundary(BSpline::natural);
                xCurve.setupSpline();
                yCurve.setData(knots, yData);
                yCurve.setBoundary(BSpline::natural);
                yCurve.setupSpline();
            } else if (dimension==3) {
                xCurve.setData(knots, xData);
                xCurve.setBoundary(BSpline::natural);
                xCurve.setupSpline();
                yCurve.setData(knots, yData);
                yCurve.setBoundary(BSpline::natural);
                yCurve.setupSpline();
                zCurve.setData(knots, zData);
                zCurve.setBoundary(BSpline::natural);
                zCurve.setupSpline();
            } else {
                throw std::invalid_argument("invalid dimension");
            }
        }
        //
        vector_type interpolate(double t){
            if (t>1||t<0) {
                throw std::invalid_argument("invalid input");
            }
            if (dimension==2) {
                vector_type point(2);
                point[0] = xCurve.interpolate(t);
                point[1] = yCurve.interpolate(t);
                return point;
            } else if (dimension==3) {
                if (cv_type==regular) {   
                    vector_type point(3);
                    point[0] = xCurve.interpolate(t);
                    point[1] = yCurve.interpolate(t);
                    point[2] = zCurve.interpolate(t);
                    return point;
                } else {
                    vector_type point = flat_to_ball(xBall.interpolate(t), yBall.interpolate(t));
                    double y_tmp = point[1];
                    double z_tmp = point[2];
                    if (theta!=0) {
                        point[1] = y_tmp*cos(theta) + z_tmp*sin(theta);
                        point[2] = -1*y_tmp*sin(theta) + z_tmp*cos(theta);
                    }
                    return point;
                }
            } else {
                throw std::invalid_argument("invalid dimension");
            }
        }
        //
        std::vector<double> ball_to_flat(double x, double y, double z){
            vector_type points;
            points.push_back(x/(1-z));
            points.push_back(y/(1-z));
            return points;
        }
        //
        std::vector<double> flat_to_ball(double X, double Y){
            vector_type points;
            double d = 1 + X*X + Y*Y;
            points.push_back(2*X/d);
            points.push_back(2*Y/d);
            points.push_back((d-2)/d);
            return points;
        }
        //
        void setupBallCurve(){
            checkData();
            cv_type = ball;
            int n = xData.size();
            while (*std::max_element(zData.begin(),zData.end())>=0.95) {
                theta += 1;
                for (int i=0; i<n; i++) {
                    double y_tmp = yData[i];
                    double z_tmp = zData[i];
                    yData[i] = y_tmp*cos(1) - z_tmp*sin(1);
                    zData[i] = y_tmp*sin(1) + z_tmp*cos(1);
                }
            }

            // project 3d ball point to 2d
            vector_type X, Y;
            for (int i=0; i<n; i++) {
                vector_type p = ball_to_flat(xData[i], yData[i], zData[i]);
                X.push_back(p[0]);
                Y.push_back(p[1]);
            }

            // generate knots for projecting points
            knots.resize(n);
            knots[0] = 0;
            if (kn_type == cumulative) {
                for (int i=1; i<n; i++) {
                    knots[i] = sqrt(pow(X[i]-X[i-1],2)+pow(Y[i]-Y[i-1],2))+knots[i-1];
                }
                for (int i=1; i<n; i++) {
                    knots[i] = knots[i]/knots[n-1];
                }
            } else {
                for (int i=1 ; i<n; i++) {
                    knots[i] = 1.0*i/n-1;
                }
            }

            // set up spline for projecting points
            xBall.setData(knots, X);
            xBall.setBoundary(BSpline::natural);
            xBall.setupSpline();
            yBall.setData(knots, Y);
            yBall.setBoundary(BSpline::natural);
            yBall.setupSpline();
        }
};

#endif